import torch, os, json
import matplotlib.pyplot as plt

# svd分解
def svd_decompose(tensor, k):
    U, S, V = torch.svd(tensor)
    U_k = U[:, :k]
    S_k = S[:k]
    V_k = V[:, :k]
    SV_k = torch.mm(torch.diag(S_k), V_k.t())
    return U_k.data.tolist(), SV_k.data.tolist()



def mask1(matrix, threshold, w=0.7):
    # 取绝对值
    matrix = torch.abs(matrix)
    # 归一化到 [0, 1]
    matrix = (matrix - matrix.min()) / (matrix.max() - matrix.min())
    # 构建 mask 张量
    mask_matrix = torch.where(matrix > threshold, torch.tensor(w), 
                              torch.tensor(1 - w))
    return mask_matrix

def mask2(matrix, threshold, w=0.7):
    
    # 取绝对值
    matrix_abs = torch.abs(matrix)
    # 展平矩阵并排序
    flattened = matrix_abs.view(-1)
    sorted_indices = torch.argsort(flattened, descending=True)
    # 计算分界点索引
    split_idx = int(threshold * len(flattened))
    # 构建 mask
    mask_flattened = torch.ones_like(flattened) * (1 - w)
    mask_flattened[sorted_indices[:split_idx]] = w
    # 恢复为原始形状
    mask_matrix = mask_flattened.view_as(matrix)
    return mask_matrix

# 加权hessian矩阵分解
def hessian_decompose(newDomainName, name, param_data, rank):
    print("正在进行加权分解", name)
    hessian_path = os.path.join('./ABSA_lowrank/hessian/{}/'.format(newDomainName), name.replace('.weight', '.json'))
    with open (hessian_path, 'r') as f:
        hessian_data = json.load(f)
    for n, matrix in hessian_data.items():
        hessian_matrix = torch.tensor(matrix)
    
    param_data_old = param_data.clone()
    # param_data = torch.exp(param_data+6)
    # 利用SVD进行初始化
    U, S, V = torch.svd(param_data)
    U_k = U[:, :rank]
    S_k = S[:rank]
    V_k = V[:, :rank]
    SV_k = torch.mm(torch.diag(S_k), V_k.t())

    # A = torch.rand_like(U_k)
    # B = torch.rand_like(SV_k) 
    A = U_k
    B = SV_k

    A.requires_grad_(True)
    B.requires_grad_(True)

    optimizer_A = torch.optim.Adam([A, B], lr=0.01)
    # optimizer_B = torch.optim.Adam([B], lr=0.01)

    # 计算mask
    # 前threshold*100%的参数设置为w
    mask_matrix = mask2(hessian_matrix, threshold=0.5, w=0.7)
    mask_matrix = mask_matrix.to('cuda:0')

    num_epochs = 1000
    # param_data = torch.exp(param_data+6)
    L = []
    for epoch in range(num_epochs):
        optimizer_A.zero_grad()
        # 重构矩阵
        M_approx = torch.matmul(A, B)
        # 计算loss
        loss = torch.norm((mask_matrix * (param_data - M_approx)))
        # loss = norm((M - M_approx))
        # 反向传播
        loss.backward()
        optimizer_A.step()
        L.append(loss.item())
        # 每10次迭代输出一次loss
        if (epoch + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item()}')
    return A, B



    

